{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tomasonjo/blogs/blob/master/llm/autogen_neo4j.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --quiet pyautogen langchain-community neo4j"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ll_0KjONDKPS",
        "outputId": "161dae79-04b6-448b-db4c-48a0ec612bf4"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m284.6/284.6 kB\u001b[0m \u001b[31m4.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.1/2.1 MB\u001b[0m \u001b[31m14.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m203.0/203.0 kB\u001b[0m \u001b[31m11.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Installing backend dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m45.5/45.5 kB\u001b[0m \u001b[31m4.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m147.8/147.8 kB\u001b[0m \u001b[31m14.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m296.7/296.7 kB\u001b[0m \u001b[31m16.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m320.7/320.7 kB\u001b[0m \u001b[31m17.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m29.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m973.5/973.5 kB\u001b[0m \u001b[31m30.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m310.2/310.2 kB\u001b[0m \u001b[31m22.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m124.3/124.3 kB\u001b[0m \u001b[31m11.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.3/49.3 kB\u001b[0m \u001b[31m4.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.0/53.0 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m142.5/142.5 kB\u001b[0m \u001b[31m4.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.6/75.6 kB\u001b[0m \u001b[31m6.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.9/77.9 kB\u001b[0m \u001b[31m6.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m4.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for neo4j (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from langchain_community.graphs import Neo4jGraph\n",
        "\n",
        "os.environ[\"NEO4J_URI\"] = 'neo4j+s://demo.neo4jlabs.com'\n",
        "os.environ[\"NEO4J_USERNAME\"] = 'recommendations'\n",
        "os.environ[\"NEO4J_PASSWORD\"] = 'recommendations'\n",
        "os.environ[\"NEO4J_DATABASE\"] = 'recommendations'\n",
        "os.environ[\"OPENAI_API_KEY\"] = 'sk-'\n",
        "\n",
        "graph = Neo4jGraph(sanitize=True)"
      ],
      "metadata": {
        "id": "7iqBa9llDuCM"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Dict\n",
        "\n",
        "def execute_query(cypher_query: str) -> Dict:\n",
        "    return graph.query(cypher_query)"
      ],
      "metadata": {
        "id": "GLMJMbj_ElE2"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "Kmu03XRHC9PE",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 104
        },
        "outputId": "f1e273f2-fb87-4784-b070-469a125d95b9"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<function __main__.execute_query(cypher_query: str) -> Dict>"
            ],
            "text/html": [
              "<div style=\"max-width:800px; border: 1px solid var(--colab-border-color);\"><style>\n",
              "      pre.function-repr-contents {\n",
              "        overflow-x: auto;\n",
              "        padding: 8px 12px;\n",
              "        max-height: 500px;\n",
              "      }\n",
              "\n",
              "      pre.function-repr-contents.function-repr-contents-collapsed {\n",
              "        cursor: pointer;\n",
              "        max-height: 100px;\n",
              "      }\n",
              "    </style>\n",
              "    <pre style=\"white-space: initial; background:\n",
              "         var(--colab-secondary-surface-color); padding: 8px 12px;\n",
              "         border-bottom: 1px solid var(--colab-border-color);\"><b>execute_query</b><br/>def execute_query(cypher_query: str) -&gt; Dict</pre><pre class=\"function-repr-contents function-repr-contents-collapsed\" style=\"\"><a class=\"filepath\" style=\"display:none\" href=\"#\">/content/&lt;ipython-input-6-3eb146711427&gt;</a>&lt;no docstring&gt;</pre></div>"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ],
      "source": [
        "from autogen import ConversableAgent\n",
        "\n",
        "# Let's first define the assistant agent that suggests tool calls.\n",
        "assistant = ConversableAgent(\n",
        "    name=\"Assistant\",\n",
        "    system_message=\"You are a helpful AI assistant. \"\n",
        "    \"You can help with answering questions about movies and actors. \"\n",
        "    \"The graph schema is \" + graph.schema +\n",
        "    \"Return 'TERMINATE' when the task is done.\",\n",
        "    llm_config={\"config_list\": [{\"model\": \"gpt-4o\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}]},\n",
        ")\n",
        "\n",
        "# The user proxy agent is used for interacting with the assistant agent\n",
        "# and executes tool calls.\n",
        "user_proxy = ConversableAgent(\n",
        "    name=\"User\",\n",
        "    llm_config=False,\n",
        "    is_termination_msg=lambda msg: msg.get(\"content\") is not None and \"TERMINATE\" in msg[\"content\"],\n",
        "    human_input_mode=\"NEVER\",\n",
        ")\n",
        "\n",
        "# Register the tool signature with the assistant agent.\n",
        "assistant.register_for_llm(name=\"graph_retrieval\", description=\"Graph query execution\")(execute_query)\n",
        "\n",
        "# Register the tool function with the user proxy agent.\n",
        "user_proxy.register_for_execution(name=\"graph_retrieval\")(execute_query)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "6aq4Is7-C9PF",
        "outputId": "f113ac7c-cf8a-4efa-8c3b-bd303354aa84",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "User (to Assistant):\n",
            "\n",
            "What is a good comedy?\n",
            "\n",
            "--------------------------------------------------------------------------------\n",
            "\n",
            ">>>>>>>> USING AUTO REPLY...\n",
            "Assistant (to User):\n",
            "\n",
            "Sure, I can recommend some good comedy movies. Would you like suggestions based on highly rated comedies, popular comedies, or something specific?\n",
            "\n",
            "--------------------------------------------------------------------------------\n",
            "User (to Assistant):\n",
            "\n",
            "\n",
            "\n",
            "--------------------------------------------------------------------------------\n",
            "\n",
            ">>>>>>>> USING AUTO REPLY...\n",
            "Assistant (to User):\n",
            "\n",
            "***** Suggested tool call (call_bb5a5uSAzpIWHFBZyO8GskjX): graph_retrieval *****\n",
            "Arguments: \n",
            "{\"cypher_query\":\"MATCH (m:Movie)-[:IN_GENRE]->(:Genre {name: 'Comedy'}) RETURN m.title AS title ORDER BY m.imdbRating DESC LIMIT 1\"}\n",
            "********************************************************************************\n",
            "\n",
            "--------------------------------------------------------------------------------\n",
            "\n",
            ">>>>>>>> EXECUTING FUNCTION graph_retrieval...\n",
            "User (to Assistant):\n",
            "\n",
            "User (to Assistant):\n",
            "\n",
            "***** Response from calling tool (call_bb5a5uSAzpIWHFBZyO8GskjX) *****\n",
            "[{\"title\": \"Ice Age: The Great Egg-Scapade\"}]\n",
            "**********************************************************************\n",
            "\n",
            "--------------------------------------------------------------------------------\n",
            "\n",
            ">>>>>>>> USING AUTO REPLY...\n",
            "Assistant (to User):\n",
            "\n",
            "A highly-rated comedy movie is \"Ice Age: The Great Egg-Scapade.\" You might enjoy watching it! If you need more recommendations or information about the movie, feel free to ask.\n",
            "\n",
            "--------------------------------------------------------------------------------\n",
            "User (to Assistant):\n",
            "\n",
            "\n",
            "\n",
            "--------------------------------------------------------------------------------\n",
            "\n",
            ">>>>>>>> USING AUTO REPLY...\n",
            "Assistant (to User):\n",
            "\n",
            "TERMINATE\n",
            "\n",
            "--------------------------------------------------------------------------------\n"
          ]
        }
      ],
      "source": [
        "chat_result = user_proxy.initiate_chat(assistant, message=\"What is a good comedy?\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "autogen",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.5"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}